{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline  \n",
    "\n",
    "import glob\n",
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "class StopWatch():\n",
    "    startTime = time.time()\n",
    "    def __init__(self):\n",
    "        self.startTime = time.time()\n",
    "    def StartTime(self):\n",
    "        self.startTime = time.time()\n",
    "    def CheckTime(self):\n",
    "        return time.time() - self.startTime\n",
    "\n",
    "    def PrintCheckTime(self, msg):\n",
    "        elapsedTime = self.CheckTime()\n",
    "        print msg + ' : ' + str(elapsedTime) + ' sec, ' + str(elapsedTime/60) + ' min'\n",
    "\n",
    "class TrainingPlot():\n",
    "\n",
    "    listXcoord = []\n",
    "    listTrainLoss = []\n",
    "    listTrainAcc = []\n",
    "    listTestLoss = []\n",
    "    listTestAcc = []\n",
    "    iterSampleCount = 0\n",
    "    iter = 0\n",
    "    sumTestAcc = 0\n",
    "    batchSize = 0\n",
    "    sampleCount = 0\n",
    "    totalIter = 0\n",
    "    avgTestAcc = 0\n",
    "    watchTotal = StopWatch()\n",
    "    watchSingle = StopWatch()\n",
    "\n",
    "    def __init__(self):\n",
    "\n",
    "        plt.ion()\n",
    "        plt.show()\n",
    "        self.watchTotal.StartTime()\n",
    "        self.watchSingle.StartTime()\n",
    "\n",
    "    def SetConfig(self, batchSize, sampleCount, totalIter):\n",
    "        self.batchSize = batchSize\n",
    "        self.sampleCount = sampleCount\n",
    "        self.totalIter = totalIter\n",
    "\n",
    "\n",
    "    def AddTest(self, trainAcc, testAcc):\n",
    "        #self.sampleCount += self.batchSize\n",
    "        self.iter += 1\n",
    "        self.listXcoord.append(self.iter)\n",
    "        self.listTrainAcc.append(trainAcc)\n",
    "        self.listTestAcc.append(testAcc)\n",
    "\n",
    "        self.sumTestAcc += testAcc\n",
    "        self.avgTestAcc = self.sumTestAcc/self.iter\n",
    "        print ('#Iter %d / %d :TrainAcc %f, TestAcc %f, AvgAcc %f' %\n",
    "               (self.iter, self.totalIter, trainAcc,testAcc,self.avgTestAcc))\n",
    "\n",
    "    def Add(self, iter, trainLoss, testLoss, trainAcc, testAcc):\n",
    "        self.sampleCount += self.batchSize\n",
    "        self.iter = iter\n",
    "        self.listXcoord.append(self.iter)\n",
    "        self.listTrainLoss.append(trainLoss)\n",
    "        self.listTrainAcc.append(trainAcc)\n",
    "        self.listTestLoss.append(testLoss)\n",
    "        self.listTestAcc.append(testAcc)\n",
    "        totalElapsedTime = self.watchTotal.CheckTime()\n",
    "        singleEapsedTime = self.watchSingle.CheckTime()\n",
    "        self.watchSingle.StartTime()\n",
    "        #print self.watch.startTime,elapsedTime\n",
    "        totalEstimatedTime = (float(self.totalIter) / (float(self.iter) + 0.0001)) * totalElapsedTime\n",
    "        print ('#Iter %d / %d : Train %f, Test %f, TrainAcc %f, TestAcc %f (1 iter %.02f sec, total %.02f min, %.02f hour remained)' %\n",
    "               (self.iter, self.totalIter,trainLoss, testLoss, trainAcc,testAcc,\n",
    "                singleEapsedTime, totalElapsedTime/60, (totalEstimatedTime - totalElapsedTime)/3600))\n",
    "\n",
    "\n",
    "    def Show(self):\n",
    "        plt.clf()\n",
    "        plt.subplot(1,2,1)\n",
    "        plt.title('Loss : %.03f' % (self.listTestLoss[-1]))\n",
    "        if (len(self.listTrainLoss) > 0):\n",
    "            plt.plot(self.listXcoord,self.listTrainLoss, color= '#%02x%02x%02x' % (58,146,204) )\n",
    "        if (len(self.listTestLoss) > 0):\n",
    "            plt.plot(self.listXcoord,self.listTestLoss, color= '#%02x%02x%02x' % (220,98,45))\n",
    "        plt.legend(['train', 'val'])\n",
    "        plt.subplot(1,2,2)\n",
    "        plt.title('Acc : %.03f' % (self.listTestAcc[-1]))\n",
    "        if (len(self.listTrainAcc) > 0):\n",
    "            plt.plot(self.listXcoord,self.listTrainAcc, color= '#%02x%02x%02x' % (58,146,204) )\n",
    "        if (len(self.listTestAcc) > 0):\n",
    "            plt.plot(self.listXcoord,self.listTestAcc, color='#%02x%02x%02x' % (34,177,76))\n",
    "        plt.ylim([0,1])\n",
    "\n",
    "        plt.draw()\n",
    "        plt.pause(0.001)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Convolutional Encoder Decoder Net\n",
    "\n",
    "Usage :\n",
    "1. Download CamVid dataset ()\n",
    "2. Run createDB once (Set following condition to 1)\n",
    "# Create DB (run once)\n",
    "if (0):\n",
    "\n",
    "3. Reset condition to 0 and run training\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "import tensorflow as tf\n",
    "# from TrainingPlot import *\n",
    "from PIL import Image\n",
    "import cPickle as pkl\n",
    "import time\n",
    "width = 128 # 320\n",
    "height = 128 # 224\n",
    "classes = 22\n",
    "kernelSize = 7\n",
    "featureSize = 64\n",
    "resumeTraining = True\n",
    "\n",
    "weights = {\n",
    "    'ce1': tf.get_variable(\"ce1\",shape= [kernelSize,kernelSize, 3, featureSize]),\n",
    "    'ce2': tf.get_variable(\"ce2\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'ce3': tf.get_variable(\"ce3\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'ce4': tf.get_variable(\"ce4\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'cd4': tf.get_variable(\"cd4\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'cd3': tf.get_variable(\"cd3\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'cd2': tf.get_variable(\"cd2\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'cd1': tf.get_variable(\"cd1\",shape= [kernelSize,kernelSize, featureSize, featureSize]),\n",
    "    'dense_inner_prod': tf.get_variable(\"dense_inner_prod\",shape= [1, 1, featureSize,classes])\n",
    "}\n",
    "\n",
    "def CreateDB(categoryName):\n",
    "    pathLoad1 = categoryName\n",
    "    pathLoad2 = categoryName + 'annot'\n",
    "    # curPath = os.path.dirname(os.path.abspath(__file__))\n",
    "    curPath = os.getcwd()\n",
    "    print curPath\n",
    "    fileList1 = glob.glob(curPath + '/' + pathLoad1 + '/*.png')\n",
    "\n",
    "    #fileList2 = glob.glob(pathLoad2 + '/*.png')\n",
    "    #plt.ion()\n",
    "    trainFile = open(categoryName + '.txt','wt')\n",
    "    count = 0\n",
    "    occupancyList = []\n",
    "    for filename in fileList1:\n",
    "        img1 = Image.open(filename)\n",
    "        #filename2 = filename.replace('_IPMImg', '_IPMLabel')\n",
    "        filename2 = curPath + '/' + pathLoad2 + '/' + os.path.basename(filename)\n",
    "        img2 = Image.open(filename2)\n",
    "\n",
    "        print >> trainFile, filename\n",
    "        print >> trainFile, filename2\n",
    "\n",
    "        #cropimg.save(pathSave + '/' + filename)\n",
    "        #plt.title(pathCity)\n",
    "        #plt.imshow(cropimg)\n",
    "        #plt.show()\n",
    "        count += 1\n",
    "\n",
    "    print ('%d data list created' % count)\n",
    "    trainFile.close()\n",
    "    return occupancyList\n",
    "\n",
    "# input : [m x h x w x c]\n",
    "def Unpooling(inputOrg, size, mask=None):\n",
    "    # m, c, h, w order\n",
    "    # print 'start unpooling'\n",
    "    # size = tf.shape(inputOrg)\n",
    "    m = size[0]\n",
    "    h = size[1]\n",
    "    w = size[2]\n",
    "    c = size[3]\n",
    "    input = tf.transpose(inputOrg, [0, 3, 1, 2])\n",
    "    # print input.get_shape()\n",
    "    x = tf.reshape(input, [-1, 1])\n",
    "    k = np.float32(np.array([1.0, 1.0]).reshape([1,-1]))\n",
    "    # k = tf.Variable([1.0, 1.0],name=\"weights\")\n",
    "    # k = tf.reshape(k,[1,-1])\n",
    "    # k = np.array(k).reshape([1, -1])\n",
    "    output = tf.matmul(x, k)\n",
    "\n",
    "    output = tf.reshape(output,[-1, c, h, w * 2])\n",
    "    # m, c, w, h\n",
    "    xx = tf.transpose(output, [0, 1, 3, 2])\n",
    "    xx = tf.reshape(xx,[-1, 1])\n",
    "    # print xx.shape\n",
    "\n",
    "    output = tf.matmul(xx, k)\n",
    "    # m, c, w, h\n",
    "    output = tf.reshape(output, [-1, c, w * 2, h * 2])\n",
    "    output = tf.transpose(output, [0, 3, 2, 1])\n",
    "    # print mask\n",
    "    outshape = tf.pack([m, h * 2, w * 2, c])\n",
    "\n",
    "    if mask != None:\n",
    "        dense_mask = tf.sparse_to_dense(mask, outshape, output, 0)\n",
    "        # print dense_mask\n",
    "        # print 'output',output\n",
    "        # print 'mask',mask\n",
    "        # print dense_mask\n",
    "            # output = tf.mul(output, mask)\n",
    "\n",
    "        return output, dense_mask\n",
    "    else:\n",
    "        return output\n",
    "\n",
    "# max pool + stride 2 transpose conv\n",
    "def Model(X, W):\n",
    "\n",
    "    # Encoder\n",
    "    encoder1 = tf.nn.conv2d(X, W['ce1'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    encoder1 = tf.nn.batch_normalization(encoder1,0.001,1.0,0,1,0.0001)\n",
    "    encoder1 = tf.nn.relu(encoder1)\n",
    "    encoder1 = tf.nn.max_pool(encoder1, ksize=[1, 2, 2, 1],\n",
    "                        strides=[1, 2, 2, 1], padding='SAME')\n",
    "    # encoder1 = tf.nn.dropout(encoder1, 0.5)\n",
    "\n",
    "    encoder2 = tf.nn.conv2d(encoder1, W['ce2'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    encoder2 = tf.nn.batch_normalization(encoder2, 0.001, 1.0, 0, 1, 0.0001)\n",
    "    encoder2 = tf.nn.relu(encoder2)\n",
    "    encoder2 = tf.nn.max_pool(encoder2, ksize=[1, 2, 2, 1],\n",
    "                          strides=[1, 2, 2, 1], padding='SAME')\n",
    "    # encoder2 = tf.nn.dropout(encoder2, 0.5)\n",
    "\n",
    "    encoder3 = tf.nn.conv2d(encoder2, W['ce3'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    encoder3 = tf.nn.batch_normalization(encoder3, 0.001, 1.0, 0, 1, 0.0001)\n",
    "    encoder3 = tf.nn.relu(encoder3)\n",
    "    encoder3 = tf.nn.max_pool(encoder3, ksize=[1, 2, 2, 1],\n",
    "                          strides=[1, 2, 2, 1], padding='SAME')\n",
    "    # encoder3 = tf.nn.dropout(encoder3, 0.5)\n",
    "\n",
    "    encoder4 = tf.nn.conv2d(encoder3, W['ce4'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    encoder4 = tf.nn.batch_normalization(encoder4, 0.001, 1.0, 0, 1, 0.0001)\n",
    "    encoder4 = tf.nn.relu(encoder4)\n",
    "    encoder4 = tf.nn.max_pool(encoder4, ksize=[1, 2, 2, 1],\n",
    "                              strides=[1, 2, 2, 1], padding='SAME')\n",
    "    # encoder4 = tf.nn.dropout(encoder4, 0.5)\n",
    "\n",
    "    # Decoder\n",
    "    decoder4 = Unpooling(encoder4, [tf.shape(X)[0], height / 16, width / 16, featureSize])\n",
    "    decoder4 = tf.nn.conv2d(decoder4, W['cd4'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    decoder4 = tf.nn.batch_normalization(decoder4, 0.001, 1.0, 0, 1, 0.0001)\n",
    "    decoder4 = tf.nn.relu(decoder4)\n",
    "    # decoder4 = tf.nn.dropout(decoder4, 0.5)\n",
    "\n",
    "    decoder3 = Unpooling(encoder3, [tf.shape(X)[0], height/8, width/8, featureSize])\n",
    "    decoder3 = tf.nn.conv2d(decoder3, W['cd3'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    decoder3 = tf.nn.batch_normalization(decoder3, 0.001, 1.0, 0, 1, 0.0001)\n",
    "    decoder3 = tf.nn.relu(decoder3)\n",
    "    # decoder3 = tf.nn.dropout(decoder3, 0.5)\n",
    "\n",
    "    decoder2 = Unpooling(decoder3, [tf.shape(X)[0], height/4, width/4, featureSize])\n",
    "    decoder2 = tf.nn.conv2d(decoder2, W['cd2'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    decoder2 = tf.nn.batch_normalization(decoder2, 0.001, 1.0, 0, 1, 0.0001)\n",
    "    decoder2 = tf.nn.relu(decoder2)\n",
    "    # decoder2 = tf.nn.dropout(decoder2, 0.5)\n",
    "\n",
    "    decoder1 = Unpooling(decoder2, [tf.shape(X)[0], height / 2, width / 2, featureSize])\n",
    "    decoder1 = tf.nn.conv2d(decoder1, W['cd1'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "    decoder1 = tf.nn.batch_normalization(decoder1, 0.001, 1.0, 0, 1.0, 0.0001)\n",
    "    decoder1 = tf.nn.relu(decoder1)\n",
    "    # decoder1 = tf.nn.dropout(decoder1, 0.5)\n",
    "\n",
    "    output = tf.nn.conv2d(decoder1, W['dense_inner_prod'], strides=[1, 1, 1, 1], padding='SAME')\n",
    "\n",
    "    # return output, mask1, mask2, mask3\n",
    "    return output\n",
    "\n",
    "\n",
    "def DenseToOneHot(labels_dense, num_classes):\n",
    "    \"\"\"Convert class labels from scalars to one-hot vectors.\"\"\"\n",
    "    num_labels = labels_dense.shape[0]\n",
    "    index_offset = np.arange(num_labels) * num_classes\n",
    "    labels_one_hot = np.zeros((num_labels, num_classes))\n",
    "    labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1\n",
    "    return labels_one_hot\n",
    "\n",
    "\n",
    "def LoadTrainingData(filename,sampleCount=None):\n",
    "    class DataSets(object):\n",
    "        pass\n",
    "\n",
    "    datalistFile = open(filename, \"rt\")\n",
    "    fileList = datalistFile.readlines()\n",
    "    # print len(fileList)\n",
    "    data = None\n",
    "    label = None\n",
    "    if sampleCount == None:\n",
    "        sampleCount = len(fileList)\n",
    "\n",
    "    for i in range(0,sampleCount,2):\n",
    "    # for i in range(0,50,2):\n",
    "        file = fileList[i].replace('\\n','')\n",
    "        # print ('%d / %d' % (i, len(fileList)))\n",
    "        img = Image.open(file)\n",
    "        img = img.resize((width, height))\n",
    "        rgb = np.array(img).reshape(1, height, width,3)\n",
    "\n",
    "        # pixels = np.concatenate((np.array(rgb[0]).flatten(),np.array(rgb[1]).flatten(),np.array(rgb[2]).flatten()),axis=0)\n",
    "        # pixels = pixels.reshape(pixels.shape[0], 1)\n",
    "        if i == 0:\n",
    "            data = rgb\n",
    "        else:\n",
    "            data = np.concatenate((data, rgb),axis=0)\n",
    "\n",
    "        # file = fileList[i * 2 + 1].replace('\\n', '')\n",
    "        # label = Image.open(file)\n",
    "\n",
    "        file = fileList[i+1].replace('\\n', '')\n",
    "        # print i,file\n",
    "        img = Image.open(file)\n",
    "        img = img.resize((width, height), Image.NEAREST)\n",
    "        labelImage = np.array(img).reshape(1, height, width,1)\n",
    "\n",
    "        if i == 0:\n",
    "            label = labelImage\n",
    "        else:\n",
    "            # print data.shape\n",
    "            label = np.concatenate((label, labelImage), axis=0)\n",
    "    labelOneHot = np.zeros((label.shape[0],label.shape[1], label.shape[2], classes))\n",
    "    for row in range(height):\n",
    "        for col in range(width):\n",
    "            single = label[:, row, col, 0]\n",
    "            # print single.shape\n",
    "            # exit(0)\n",
    "            # print index\n",
    "            oneHot = DenseToOneHot(single, classes)\n",
    "            labelOneHot[:, row, col, :] = oneHot\n",
    "    # for i in range(22):\n",
    "    #     plt.subplot(1,2,1)\n",
    "    #     plt.imshow(data[0,:,:,:].reshape(height,width,3))\n",
    "    #     plt.subplot(1, 2, 2)\n",
    "    #     plt.imshow(labelOneHot[0,:,:,i].reshape(height,width))\n",
    "    #     plt.show()\n",
    "    return [data.astype(np.float32)/255, label, labelOneHot.astype(np.float32)]\n",
    "\n",
    "def ShowDebuggingPlot():\n",
    "    global batchData, batchLabel\n",
    "    index = np.random.randint(trainData.shape[0])\n",
    "    batchData = trainData[index:index+1]\n",
    "    batchLabel = trainLabelOneHot[index:index+1]\n",
    "    predMaxOut = sess.run(predMax, feed_dict={x: batchData, y: batchLabel})\n",
    "    yMaxOut = sess.run(yMax, feed_dict={x: batchData, y: batchLabel})\n",
    "    # for i in range(22):\n",
    "    # show predicted image\n",
    "    plt.figure(2)\n",
    "    plt.clf()\n",
    "    plt.subplot(2, 2, 1)\n",
    "    plt.title('Input')\n",
    "    img = trainData[index, :, :, :].reshape(height, width, 3)\n",
    "    # plt.imshow(img * 255)\n",
    "    plt.imshow(img)\n",
    "    plt.subplot(2, 2, 2)\n",
    "    plt.title('Ground truth')\n",
    "    img = yMaxOut[0, :, :].reshape(height, width)\n",
    "    plt.imshow(img)\n",
    "    plt.subplot(2, 2, 3)\n",
    "    plt.title('[Training] Prediction')\n",
    "    plt.imshow(predMaxOut[0, :, :].reshape(height, width))\n",
    "    plt.subplot(2, 2, 4)\n",
    "    plt.title('Error')\n",
    "    plt.imshow(img - predMaxOut[0, :, :].reshape(height, width))\n",
    "    plt.show()\n",
    "    \n",
    "def ShowValPlot():\n",
    "    global batchData, batchLabel\n",
    "    index = np.random.randint(valData.shape[0])\n",
    "    batchData = valData[index:index+1]\n",
    "    batchLabel = valLabelOneHot[index:index+1]\n",
    "    predMaxOut = sess.run(predMax, feed_dict={x: batchData, y: batchLabel})\n",
    "    yMaxOut = sess.run(yMax, feed_dict={x: batchData, y: batchLabel})\n",
    "    # for i in range(22):\n",
    "    # show predicted image\n",
    "    plt.figure(3)\n",
    "    plt.clf()\n",
    "    plt.subplot(2, 2, 1)\n",
    "    plt.title('Input')\n",
    "    img = valData[index, :, :, :].reshape(height, width, 3)\n",
    "    # plt.imshow(img * 255)\n",
    "    plt.imshow(img)\n",
    "    plt.subplot(2, 2, 2)\n",
    "    plt.title('Ground truth')\n",
    "    img = yMaxOut[0, :, :].reshape(height, width)\n",
    "    plt.imshow(img)\n",
    "    plt.subplot(2, 2, 3)\n",
    "    plt.title('[Validation] Prediction')\n",
    "    plt.imshow(predMaxOut[0, :, :].reshape(height, width))\n",
    "    plt.subplot(2, 2, 4)\n",
    "    plt.title('Error')\n",
    "    plt.imshow(img - predMaxOut[0, :, :].reshape(height, width))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/enginius/github/tensorflow-101/notebooks\n",
      "367 data list created\n",
      "/home/enginius/github/tensorflow-101/notebooks\n",
      "101 data list created\n",
      "Start data loading\n",
      "Finished in 0 sec\n",
      "Strart training..\n"
     ]
    }
   ],
   "source": [
    "# main body\n",
    "\n",
    "# Create DB (run once)\n",
    "camvidpath = 'data/seg/SegNet-Tutorial-master/CamVid/'\n",
    "if (1):\n",
    "    CreateDB(camvidpath + 'train')\n",
    "    CreateDB(camvidpath + 'val')\n",
    "    \n",
    "startTime = time.time()\n",
    "print('Start data loading')\n",
    "sampleCount = 20\n",
    "trainData, trainLabel, trainLabelOneHot = LoadTrainingData(camvidpath + 'train.txt',sampleCount)\n",
    "valData, valLabel, valLabelOneHot = LoadTrainingData(camvidpath + 'val.txt',sampleCount/2)\n",
    "\n",
    "print('Finished in %d sec' % (time.time() - startTime))\n",
    "\n",
    "# Define functions\n",
    "x = tf.placeholder(tf.float32, [None, height,width,3])\n",
    "y = tf.placeholder(tf.float32, [None, height,width,classes])\n",
    "\n",
    "# input = tf.reshape(x, shape=[-1, 224, 320, 3])\n",
    "# output : m x height x width x classes\n",
    "\n",
    "# pred, mask1, mask2, mask3 = Model(x, weights)\n",
    "pred = Model(x, weights)\n",
    "\n",
    "linearizePred = tf.reshape(pred,shape=[-1,classes])\n",
    "linearizeY = tf.reshape(y,shape=[-1,classes])\n",
    "cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(linearizePred, linearizeY))\n",
    "\n",
    "\n",
    "# accuracy\n",
    "# print yNumber, tf.argmax(pred, 3)\n",
    "predMax = tf.argmax(pred, 3)\n",
    "yMax = tf.argmax(y, 3)\n",
    "correct_pred = tf.equal(tf.argmax(y,3), tf.argmax(pred, 3)) # Count correct predictions\n",
    "\n",
    "\n",
    "acc_op = tf.reduce_mean(tf.cast(correct_pred, \"float\")) # Cast boolean to float to average\n",
    "\n",
    "learning_rate = 0.0001\n",
    "\n",
    "optm = tf.train.AdamOptimizer(learning_rate).minimize(cost)\n",
    "\n",
    "# Fit all training data\n",
    "batch_size = 2\n",
    "n_epochs = 10000\n",
    "\n",
    "print(\"Strart training..\")\n",
    "\n",
    "trainingPlot = TrainingPlot()\n",
    "trainingPlot.SetConfig(batch_size, 500, n_epochs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10, 128, 128, 3)\n",
      "(10, 128, 128, 1)\n",
      "(10, 128, 128, 22)\n"
     ]
    }
   ],
   "source": [
    "print (trainData.shape)\n",
    "print (trainLabel.shape)\n",
    "print (trainLabelOneHot.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\n",
    "# you need to initialize all variables\n",
    "with tf.Session() as sess:\n",
    "    # you need to initialize all variables\n",
    "    tf.initialize_all_variables().run()\n",
    "\n",
    "    # create a log writer. run 'tensorboard --logdir=./logs/nn_logs'\n",
    "    # if (resumeTraining):\n",
    "    saver = tf.train.Saver()\n",
    "    checkpoint = tf.train.latest_checkpoint(\"nets/deconv\")\n",
    "    if resumeTraining == False:\n",
    "        print \"Start from scratch\"\n",
    "    elif  checkpoint:\n",
    "        print \"Restoring from checkpoint\", checkpoint\n",
    "        saver.restore(sess, checkpoint)\n",
    "    else:\n",
    "        print \"Couldn't find checkpoint to restore from. Starting over.\"\n",
    "\n",
    "\n",
    "    for epoch_i in range(n_epochs):\n",
    "        trainLoss = []\n",
    "        trainAcc = []\n",
    "        for start, end in zip(range(0, len(trainData), batch_size), range(batch_size, len(trainData), batch_size)):\n",
    "\n",
    "            batchData = trainData[start:end]\n",
    "            batchLabel = trainLabelOneHot[start:end]\n",
    "\n",
    "            sess.run(optm, feed_dict={x: batchData, y: batchLabel})\n",
    "            trainLoss.append(sess.run(cost, feed_dict={x: batchData, y: batchLabel}))\n",
    "            trainAcc.append(sess.run(acc_op, feed_dict={x: batchData, y: batchLabel}))\n",
    "\n",
    "        trainLoss = np.mean(trainLoss)\n",
    "        trainAcc = np.mean(trainAcc)\n",
    "\n",
    "        # run validation\n",
    "        valLoss = sess.run(cost, feed_dict={x: valData, y: valLabelOneHot})\n",
    "        valAcc = sess.run(acc_op, feed_dict={x: valData, y: valLabelOneHot})\n",
    "\n",
    "        # trainingPlot.Add(epoch_i, trainLoss, valLoss, trainAcc, valAcc)\n",
    "        # plt.figure(1)\n",
    "        # trainingPlot.Show()\n",
    "\n",
    "        # save snapshot\n",
    "        if resumeTraining and epoch_i % 1000 == 0:\n",
    "            print \"training on image #%d\" % epoch_i\n",
    "            saver.save(sess, 'nets/deconv/progress', global_step=epoch_i)\n",
    "\n",
    "        # show debugging image\n",
    "        if epoch_i % 500 == 0:\n",
    "            trainingPlot.Add(epoch_i, trainLoss, valLoss, trainAcc, valAcc)\n",
    "            plt.figure(1)\n",
    "            trainingPlot.Show()\n",
    "            # Debug plot\n",
    "            ShowDebuggingPlot()\n",
    "            # Plot on validation set\n",
    "            ShowValPlot()\n",
    "\n",
    "        # print(epoch_i, \"/\", n_epochs, loss)\n",
    "\n",
    "print(\"Training done. \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
